# 本脚本未完成编辑
import math

import numpy as np
import pandas as pd

df = pd.read_excel('../data/创业板指研究.xlsx', sheet_name='Sheet1')


# 策略选择函数，共有5×3种（机会判断依据、指数最长持有时间）
def cl_xz(x):
    cy_sj, pd_sj, shj = 0, 0, False  # 分别是“持有时间”“判断时间”“是否衰减”
    # “持有时间”是一直不满足卖出条件时，为保持资金流动性，不得不卖出时，距离买入日期的时间，即单次最长持有时间
    # “判断时间”是判断是否购买时的依据，其数值为买入时，计算前面指数涨跌幅的天数
    # “是否衰减”是指是计算前几天涨跌幅时，是直接叠加涨跌幅，还是降低距离今天更远的时间的涨跌幅的权重
    if x % 3 == 0:
        cy_sj = 15
    elif x % 3 == 1:
        cy_sj = 45
    else:
        cy_sj = 75
    if x // 3 == 0:
        pd_sj = 1
    elif x // 3 == 1:
        pd_sj = 3
    elif x // 3 == 2:
        pd_sj = 3
        shj = True
    elif x // 3 == 3:
        pd_sj = 5
    else:
        pd_sj = 5
        shj = True
    return [cy_sj, pd_sj, shj]


def pd_yj(pd_sj, shj):  # 计算判断依据'yj'列
    if shj == True:
        if pd_sj == 3:
            df['yj'] = (df['涨幅(%)'] + 0.8 * df['涨幅(%)'].shift(1) + 0.5 * df['涨幅(%)'].shift(2))
            df.loc[0:2, 'yj'] = 0  # 将前三行设为0，避免出现空值，后面不再注释
        else:
            df['yj'] = (df['涨幅(%)'] + 0.8 * df['涨幅(%)'].shift(1) + 0.6 * df['涨幅(%)'].shift(2)
                        + 0.4 * df['涨幅(%)'].shift(3) + 0.2 * df['涨幅(%)'].shift(4))
            df.loc[0:4, 'yj'] = 0
    else:
        if pd_sj == 1:
            df['yj'] = df['涨幅(%)']
            df.loc[0, 'yj'] = 0
        if pd_sj == 3:
            df['yj'] = (df['涨幅(%)'] + df['涨幅(%)'].shift(1) + df['涨幅(%)'].shift(2))
            df.loc[0:2, 'yj'] = 0
        if pd_sj == 5:
            df['yj'] = (df['涨幅(%)'] + df['涨幅(%)'].shift(1) + df['涨幅(%)'].shift(2)
                        + df['涨幅(%)'].shift(3) + df['涨幅(%)'].shift(4))
            df.loc[0:4, 'yj'] = 0


def main(cy_sj, pd_sj, shj, tj, qs_year, js_year):  # 主要逻辑函数，根据规则交易，并计算收益率
    # 参数tj为二元列表，[买入时yj的下限,买入时yj的上限]
    # 参数qs_year表示研究起始年份,js_year表示研究结束年份
    pd_yj(pd_sj, shj)
    df['worst'] = (df['收盘价'].shift(-cy_sj) - df['收盘价']) / df['收盘价'] * 100
    df.loc[df['worst'].isna().tail(cy_sj).index, 'worst'] = 0  # 清除最后的缺失值
    yj = df['yj'].values  # 投资判断依据（表列转化为列表）
    spj = df['收盘价'].values  # 指数收盘价
    year, month, day = df[['year', 'month', 'day']].to_numpy().T  # 将日期转换为列表，后续会用到
    # 最关键的一行，将不满足条件的日子从投资列表中剔除
    btr = [i for i, value in enumerate(yj) if value < tj[0] or value > tj[1]]  # 表示不投入的交易日
    n = 0  # 初始化——投资失败次数
    days = []  # 初始化——单个投资周期天数
    s = []  # 初始化——每次投资收益率
    tz_date = []  # 初始化——每次投资的日期
    yj_m = [100] * 4 + [0] * 3 + [-2] * (cy_sj - 7)  # 卖出时判断依据的涨幅（%）,可更改
    for i in range(len(spj) - cy_sj):
        if (i in btr) or (year[i] < qs_year or year[i] > js_year):  # 去除早于起始研究年限或晚于结束研究年限的，以及按规则不投资的数据
            continue
        for j in range(1, cy_sj + 1):
            if spj[i + j] > spj[i] * (1 + yj_m[j - 1] / 100) and spj[i + j] > spj[i + j - 1]:  # 是否卖出依据浮盈是否超过yj_m对应预设确定
                if j <= 4:
                    s.append((spj[i + j] * 0.985 / spj[i] - 1) * 100)  # 在前4个交易日内卖出会收1.5%的手续费
                else:
                    s.append((spj[i + j] / spj[i] - 1) * 100)  # 一周后卖出无手续费
                tz_date.append([year[i], month[i], day[i]])  # 记录盈利的投资日期
                days.append(j)
                break
            if j == cy_sj:
                s.append((spj[i + cy_sj] / spj[i] - 1) * 100)  # 限定的交易日内未能盈利，按最后一天的价格卖出
                tz_date.append([year[i], month[i], day[i]])  # 记录失败的投资日期
                n += 1
                days.append(cy_sj)
    print(f'持有时间：{cy_sj}\t判断依据：{pd_sj} {tj}\t是否衰减：{shj}\t投资次数：{len(s)}')
    print(f'胜率：{(1 - n / len(s)) * 100:.2f}%')
    print(f'平均收益率（算数平均）：{np.mean(s):.2f}%')
    jh_pj = (math.exp(np.mean([math.log(s[i] / 100 + 1, math.e) for i in range(len(s))])) - 1) * 100
    print(f'平均收益率（几何平均）：{jh_pj:.2f}%')
    print(f'理想情况总收益率：顺投：{((jh_pj / 100 + 1) ** len(s) - 1) * 100:.2f}% 等额：{np.sum(s):.2f}%')
    print(f'平均周转天数：{np.mean(days):.2f}')

    # 整理出决定要投资的日子的数据表
    tz_year = [date[0] for date in tz_date]
    tz_month = [date[1] for date in tz_date]
    tz_day = [date[2] for date in tz_date]
    data = [
        tz_year,  # 年份
        tz_month,  # 月份
        tz_day,  # 日期
        s,  # 收益率
        days  # 投资持有期
    ]
    lm = ['年', '月', '日', '收益率', '持有时长']  # 表列名字
    data = list(zip(*data))  # 转置表格
    yjy = pd.DataFrame(data, columns=lm)  # 每次交易的示例
    yjy.to_excel(f'../output/创业板指投资示例：{cy_sj}天+依据{pd_sj}天数据+{shj * '非'}线性.xlsx', sheet_name='Sheet1',
                 index=False)


if __name__ == '__main__':
    qs_year = 2012  # 起始年份（含该年）
    js_year = 2025  # 结束年份（含该年）
    for i in range(15):
        cy_sj, pd_sj, shj = cl_xz(i)
        pd_sj_1 = (pd_sj - 3) * (pd_sj - 5) / 8
        pd_sj_3 = (pd_sj - 1) * (pd_sj - 5) / (-4)
        pd_sj_5 = (pd_sj - 1) * (pd_sj - 3) / 8
        if shj == True:
            tj = [2 * pd_sj_1 + 2.5 * pd_sj_3 + 3 * pd_sj_5,
                  4.5 * pd_sj_1 + 6 * pd_sj_3 + 7.5 * pd_sj_5]
        else:
            tj = [2 * pd_sj_1 + 3 * pd_sj_3 + 4 * pd_sj_5,
                  4.5 * pd_sj_1 + 7.5 * pd_sj_3 + 10.5 * pd_sj_5]
        main(cy_sj, pd_sj, shj, tj, qs_year, js_year)
